import os
from datetime import datetime

from langchain_community.tools import TavilySearchResults
from openai import OpenAI

# 配置信息集中管理
CONFIG = {
    "tavily_api_key": "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",
    "model_endpoint": "https://api-inference.modelscope.cn/v1/",
    "model_token": "b011f8fd-99a5-4633-8bf5-fef46cf7bd90",
    "model_name": "deepseek-ai/DeepSeek-R1-Distill-Qwen-7B"
}
os.environ["TAVILY_API_KEY"] = CONFIG["tavily_api_key"]

client = OpenAI(
    base_url=CONFIG["model_endpoint"],
    api_key=CONFIG["model_token"],
)

web_tool = TavilySearchResults(
    max_results=5,
    search_depth="advanced",
    include_answer=True,
    include_raw_content=True,
    include_images=False
)


def web_search(query: str, max_retry: int = 1) -> str:
    """获取网络搜索结果并进行格式化"""
    for attempt in range(max_retry):
        try:
            result = web_tool.invoke({"query": query})
            formatted_result = ""
            for i, item in enumerate(result, start=1):
                page_content = item['content']
                formatted_result += (
                    f"[webpage {i} begin]\n"
                    f"URL: {item['url']}\n"
                    f"Content: {page_content[:2000]}...\n"  # 限制内容长度
                    f"[webpage {i} end]\n"
                )
            return formatted_result
        except Exception as e:
            print(e)
            if attempt == max_retry - 1:
                return f"网络搜索暂时不可用，错误信息：{str(e)}"
            continue


def web_search_prompt(search_results: str, query: str) -> str:
    """构建系统提示模板"""
    current_time = datetime.now().strftime('%Y-%m-%d %H:%M:%S')

    return f"""
# 以下内容是基于用户发送的消息的搜索结果:{search_results}

在我给你的搜索结果中，每个结果都是[webpage X begin]...[webpage X end]格式的，X代表每篇文章的数字索引。
请在适当的情况下在句子末尾引用上下文。请按照引用编号[citation:X]的格式在答案中对应部分引用上下文。如果一句话源自多个上下文，请列出所有相关的引用编号，例如[citation:3][citation:5]，切记不要将引用集中在最后返回引用编号，而是在答案对应部分列出。
在回答时，请注意以下几点：
- 今天是{current_time}。
- 并非搜索结果的所有内容都与用户的问题密切相关，你需要结合问题，对搜索结果进行甄别、筛选。
- 对于列举类的问题（如列举所有航班信息），尽量将答案控制在10个要点以内，并告诉用户可以查看搜索来源、获得完整信息。优先提供信息完整、最相关的列举项；如非必要，不要主动告诉用户搜索结果未提供的内容。
- 对于创作类的问题（如写论文），请务必在正文的段落中引用对应的参考编号，例如[citation:3][citation:5]，不能只在文章末尾引用。你需要解读并概括用户的题目要求，选择合适的格式，充分利用搜索结果并抽取重要信息，生成符合用户要求、极具思想深度、富有创造力与专业性的答案。你的创作篇幅需要尽可能延长，对于每一个要点的论述要推测用户的意图，给出尽可能多角度的回答要点，且务必信息量大、论述详尽。
- 如果回答很长，请尽量结构化、分段落总结。如果需要分点作答，尽量控制在5个点以内，并合并相关的内容。
- 对于客观类的问答，如果问题的答案非常简短，可以适当补充一到两句相关信息，以丰富内容。
- 你需要根据用户要求和回答内容选择合适、美观的回答格式，确保可读性强。
- 你的回答应该综合多个相关网页来回答，不能重复引用一个网页。
- 除非用户要求，否则你回答的语言需要和用户提问的语言保持一致。
# 用户消息为：{query}"""


def generate_answer(query: str) -> None:
    """生成答案主流程"""

    # 判断是否需要联网
    web_flag = check_web_search(user_query)
    if web_flag:
        # 获取网络搜索结果
        print("\n=== 正在搜索最新信息 ===")
        search_results = web_search(query)
        # 构建提示词
        user_prompt = web_search_prompt(search_results, query)
    else:
        # 只能回答
        print("\n=== 直接回答 ===")
        user_prompt = user_query
    print(user_prompt)
    try:
        response = client.chat.completions.create(
            model=CONFIG["model_name"],
            messages=[
                {"role": "user", "content": user_prompt}
            ],
            stream=True,
            temperature=0.6,  # 降低随机性
            max_tokens=8000
        )
        print("\n\n=== 生成回答 ===")
        done_reasoning = False
        for chunk in response:
            reasoning_chunk = chunk.choices[0].delta.reasoning_content
            answer_chunk = chunk.choices[0].delta.content
            if reasoning_chunk != '':
                print(reasoning_chunk, end='', flush=True)
            elif answer_chunk != '':
                if not done_reasoning:
                    print('\n\n === Final Answer ===\n')
                    done_reasoning = True
                print(answer_chunk, end='', flush=True)

    except Exception as e:
        print(e)
        print(f"\n生成回答时发生错误：{str(e)}")


def check_web_search(query: str) -> bool:
    current_time = datetime.now().strftime('%Y-%m-%d %H:%M:%S')
    system_prompt = f"""
    # Role：智能检索决策引擎，只能返回yes或者no
    ## 当前时间
    {current_time}
    
    ## 判断流程
    请按以下步骤判断是否需要联网：
    1. 优先检查[必须联网]条件
    2. 然后检查[不需联网]条件
    3. 最后执行默认规则
    
    ## 判断规则
    ### 必须联网（返回yes）
    - 包含以下时间敏感词：今天、现在、实时、最新、刚刚、近期、此刻、最近X天（X≤7）
    - 涉及以下实时信息类型：
      * 天气/新闻/股价/赛事/交通/热搜
      * 政策法规/医学指南/学术前沿（近2年）
      * 预测类信息（如股市预测）
    - 问题含明确时间范围（示例：2023年后、最近三个月、>2024-01-01）
    - 用户明确要求更新/最新信息
    - 需要时效性知识（如：当前国家领导人）
    
    ### 不需联网（返回no）
    - 数学计算/逻辑推理/代码调试
    - 文本处理（翻译/润色/摘要）
    - 通用知识（历史事件/经典理论/教科书内容）
    - 问题包含本地数据标识（如：根据本地文档）
    - 明确的历史数据查询（示例：2022年GDP）
    
    ### 默认规则
    当无法明确判断时：
    - 问题含时间状语 ⇒ yes
    - 其他情况 ⇒ yes
    """

    response = client.chat.completions.create(
        model=CONFIG["model_name"],
        temperature=0,
        messages=[
            {"role": "system", "content": system_prompt},
            {"role": "user", "content": f"用户问题：{query}"}
        ]
    )
    print(response.choices[0].message.reasoning_content)
    print(response.choices[0].message.content)
    return 'yes' == response.choices[0].message.content


if __name__ == "__main__":
    user_query = "2025年春运情况"
    generate_answer(user_query)
